import torch
from ultralytics import YOLO

# 根据项目实际情况导入 Detect 模块、ChannelAttention 和 ImprovedC2f_AMHA
# 以下路径可能需要根据实际项目调整
from ultralytics.nn.modules import Detect 
from ultralytics.nn.modules.MSAD import ChannelAttention, AMHA
from ultralytics.nn.modules.omniFusion import ImprovedC2f_AMHA

def print_model_layers(model):
    """
    遍历模型所有子模块，并打印出名称、类型。
    如果模块属于 Detect 类型，则输出其 cv2、cv3 和 dfl 的详细信息；
    如果模块是 ChannelAttention 或 ImprovedC2f_AMHA，也进行单独标注。
    """
    for name, module in model.named_modules():
        print(f"Layer name: {name}, type: {module.__class__.__name__}")
        
        # 如果是 Detect 模块，打印其 cv2, cv3 及 dfl 结构
        if isinstance(module, Detect):
            print("  -- Detect 模块内部结构:")
            try:
                # cv2 通常为一个列表或 ModuleList，其中每个元素可能为一个元组/list
                if hasattr(module, "cv2"):
                    print("    cv2:")
                    for i, item in enumerate(module.cv2):
                        print(f"      cv2[{i}]: {item}")
                        if isinstance(item, (list, tuple)) and len(item) > 2:
                            print(f"        cv2[{i}][2]: {item[2]}")
                if hasattr(module, "cv3"):
                    print("    cv3:")
                    for i, item in enumerate(module.cv3):
                        print(f"      cv3[{i}]: {item}")
                        if isinstance(item, (list, tuple)) and len(item) > 2:
                            print(f"        cv3[{i}][2]: {item[2]}")
                if hasattr(module, "dfl"):
                    print("    dfl:")
                    print(f"      {module.dfl}")
            except Exception as e:
                print("    Error printing Detect details:", e)
        
        # 如果是 ChannelAttention 模块
        if isinstance(module, ChannelAttention):
            print("  -- This is a ChannelAttention module.")
        
        # 如果是 ImprovedC2f_AMHA 模块
        if isinstance(module, ImprovedC2f_AMHA):
            print("  -- This is an ImprovedC2f_AMHA module.")

if __name__ == '__main__':
    # 修改此处的 model_path 为你实际的模型权重路径
    model_path = "runs/train/exp/weights/best.pt"
    # 加载模型（YOLO 类内部包含训练流程，但我们只用来观察网络结构）
    model = YOLO(model_path)
    
    # 输出模型所有层的名称和详细信息
    print_model_layers(model)
